package net.tp.struct.unionfind;

import java.util.Arrays;

import static java.lang.Math.*;

/**
 *
 * @author Trung Phan
 *
 */
public class UnionFind {

    private int[] a;
    
    private int group;

    /**
     * Construct the union find structure with N nodes.
     *
     * @param N number of nodes
     */
    public UnionFind(int N) {
        this.a = new int[N];
        Arrays.fill(a, -1);
        this.group = N;
    }

    /**
     * Union two nodes.
     *
     * @param i node i
     * @param j node j
     * @time $O(log n)$ worst case, $O(1)$ with high probability.
     */
    public void union(int i, int j) {
        if (i < 0 || i >= a.length || j < 0 || j >= a.length)
            throw new IndexOutOfBoundsException();
        
        if (i == j) return;
    
        int p = root(i);
        int q = root(j);

        if (p == q) return;

        this.group--;

        if (abs(a[p]) >= abs(a[q])) {
            a[p] += a[q];
            a[q] = p;
        }
        else {
            a[q] += a[p];
            a[p] = q;
        }
    }

    /**
     * Find the root for node v.
     * Root is the node that has value less than 0. The absolute value of the root is the
     * size of the subtree rooted at that root.
     *
     * @param v node v
     * @return the root node
     * @time $O(log n)$ worst case, $O(1)$ with high probability.
     */
    public int root(int v) {
        int root = v;
        for (; a[root] >= 0; root = a[root]);

        compressPath(v, root);

        return root;
    }

    /**
     * Make root the parent of all the nodes along the path v -> root.
     * This is the trick to make root() $O(log n)$ most of the time.
     */
    private void compressPath(int v, int root) {
        while (v != root && a[v] != root) {
            int u = a[v];
            a[v] = root;
            v = u;
        }
    }

    /**
     * Check if node i and j are connected.
     *
     * @param i node i
     * @param j node j
     * @return true if node i and j are connected
     * @time $O(log n)$ worst case, $O(1)$ with high probability
     */
    public boolean connected(int i, int j) {
        return root(i) == root(j);
    }

    /**
     * Number of groups.
     *
     * @return number of groups
     */
    public int group() {
        return this.group;
    }
}
